from torch.utils.data.sampler import BatchSampler
import torch
from easycore.common.registry import Registry
from .sampler import InfiniteSampler, InfiniteInstanceBalanceSampler, InfiniteCategoryBalanceSampler, InfiniteDatasetTypeBalanceSampler
from .dataset import build_dataset
from .collate import default_collate
from .data_mapper import MapDataset
from .yolo_dataset import YoloDataset

LOADER_REGISTRY = Registry("data loader")

def build_train_loader(cfg, mapper=None):
    return LOADER_REGISTRY.get(cfg.LOADER.TRAIN)(cfg, mapper=mapper)

def build_test_loader(cfg, mapper=None):
    return LOADER_REGISTRY.get(cfg.LOADER.TEST)(cfg, mapper=mapper)

@LOADER_REGISTRY.register()
def build_default_train_loader(cfg, mapper=None):
    batch_size = cfg.LOADER.BATCH_SIZE
    dataset_dicts = build_dataset(*cfg.DATASETS.TRAIN)

    sampler = InfiniteSampler(len(dataset_dicts))
    sampler = BatchSampler(sampler, batch_size, drop_last=True)

    if mapper is None:
        dataset = dataset_dicts
    else:
        dataset = MapDataset(dataset_dicts, mapper)

    data_loader = torch.utils.data.DataLoader(
        dataset,
        batch_sampler=sampler,
        collate_fn=default_collate,
        num_workers=cfg.LOADER.NUM_WORKERS,
    )
    return data_loader

@LOADER_REGISTRY.register()
def build_instance_balance_train_loader(cfg, mapper=None):
    batch_size = cfg.LOADER.BATCH_SIZE
    dataset_dicts = build_dataset(*cfg.DATASETS.TRAIN)

    sampler = InfiniteInstanceBalanceSampler(dataset_dicts, cfg.MODEL.NUM_CLASSES)
    sampler = BatchSampler(sampler, batch_size, drop_last=True)

    if mapper is None:
        dataset = dataset_dicts
    else:
        dataset = MapDataset(dataset_dicts, mapper)

    data_loader = torch.utils.data.DataLoader(
        dataset,
        batch_sampler=sampler,
        collate_fn=default_collate,
        num_workers=cfg.LOADER.NUM_WORKERS,
    )
    return data_loader

@LOADER_REGISTRY.register()
def build_category_balance_train_loader(cfg, mapper=None):
    batch_size = cfg.LOADER.BATCH_SIZE
    dataset_dicts = build_dataset(*cfg.DATASETS.TRAIN)

    sampler = InfiniteCategoryBalanceSampler(dataset_dicts, cfg.MODEL.NUM_CLASSES)
    sampler = BatchSampler(sampler, batch_size, drop_last=True)

    if mapper is None:
        dataset = dataset_dicts
    else:
        dataset = MapDataset(dataset_dicts, mapper)

    data_loader = torch.utils.data.DataLoader(
        dataset,
        batch_sampler=sampler,
        collate_fn=default_collate,
        num_workers=cfg.LOADER.NUM_WORKERS,
    )
    return data_loader

@LOADER_REGISTRY.register()
def build_category_balance_yolo_train_loader(cfg, mapper=None):
    batch_size = cfg.LOADER.BATCH_SIZE
    dataset_dicts = build_dataset(*cfg.DATASETS.TRAIN)

    sampler = InfiniteCategoryBalanceSampler(dataset_dicts, cfg.MODEL.NUM_CLASSES)
    sampler = BatchSampler(sampler, batch_size, drop_last=True)

    if mapper is None:
        dataset = dataset_dicts
    else:
        dataset = YoloDataset(dataset_dicts, mapper, train=True)

    data_loader = torch.utils.data.DataLoader(
        dataset,
        batch_sampler=sampler,
        collate_fn=default_collate,
        num_workers=cfg.LOADER.NUM_WORKERS,
    )
    return data_loader

@LOADER_REGISTRY.register()
def build_dataset_type_balance_train_loader(cfg, mapper=None):
    batch_size = cfg.LOADER.BATCH_SIZE
    dataset_dicts = build_dataset(*cfg.DATASETS.TRAIN)

    sampler = InfiniteDatasetTypeBalanceSampler(dataset_dicts)
    sampler = BatchSampler(sampler, batch_size, drop_last=True)

    if mapper is None:
        dataset = dataset_dicts
    else:
        dataset = MapDataset(dataset_dicts, mapper)

    data_loader = torch.utils.data.DataLoader(
        dataset,
        batch_sampler=sampler,
        collate_fn=default_collate,
        num_workers=cfg.LOADER.NUM_WORKERS,
    )
    return data_loader

@LOADER_REGISTRY.register()
def build_default_test_loader(cfg, mapper=None):
    batch_size = cfg.LOADER.BATCH_SIZE
    dataset_dicts = build_dataset(*cfg.DATASETS.TEST)

    if mapper is None:
        dataset = dataset_dicts
    else:
        dataset = MapDataset(dataset_dicts, mapper)

    data_loader = torch.utils.data.DataLoader(
        dataset,
        batch_size = batch_size,
        shuffle = False,
        collate_fn = default_collate,
    )
    return data_loader
